//
// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
//

#include "hnswlib/hnswlib.h"
#include "nlohmann/json.hpp"
#include <fstream>
#include <unordered_map>
#include <iostream>
#include <vector>
#include <string>
#include "model.h"
#include <turbo/flags/flag.h>
#include <krpc/utility/logging.h>
#include <krpc/rpc/server.h>
#include <krpc/rpc/restful.h>
#include <krpc/json2pb/pb_to_json.h>
#include <krpc/rpc/restful_service.h>
#include "version.h"


std::unordered_map<std::string, hnswlib::HierarchicalNSW<float>*> indices;
std::unordered_map<std::string, nlohmann::json> index_settings;

void write_index_to_disk(const std::string &index_name) {
    auto *index = indices[index_name];
    index->saveIndex("indices/" + index_name + ".bin");

    std::ofstream settings_file("indices/" + index_name + ".json");
    settings_file << index_settings[index_name].dump();
}

void read_index_from_disk(const std::string &index_name) {
    std::ifstream settings_file("indices/" + index_name + ".json");
    nlohmann::json index_state;
    settings_file >> index_state;

    int dim = index_state["dimension"];
    std::string space = index_state["space_type"];
    int ef_construction = index_state["ef_construction"];
    int M = index_state["M"];

    hnswlib::SpaceInterface<float>* metric_space = (space == "IP")
                                                   ? static_cast<hnswlib::SpaceInterface<float>*>(new hnswlib::InnerProductSpace(dim))
                                                   : static_cast<hnswlib::SpaceInterface<float>*>(new hnswlib::L2Space(dim));

    auto *index = new hnswlib::HierarchicalNSW<float>(
            metric_space,
            10000,
            M,
            ef_construction,
            42,
            true
    );
    std::string index_path = "indices/" + index_name + ".bin";
    index->loadIndex(index_path, metric_space, 10000);

    indices[index_name] = index;
    index_settings[index_name] = index_state;
}

static void set_response_ok(krpc::RestfulResponse*res) {
    res->set_status_code(200);
    res->set_content_type("text/plain");
    nlohmann::json obj;
    obj["code"] = 0;
    obj["message"] = "ok";
    res->set_body(obj.dump());
}

static void set_response_version(krpc::RestfulResponse*res) {
    res->set_status_code(200);
    res->set_content_type("text/plain");
    nlohmann::json obj;
    obj["name"] = "rigel";
    obj["version"] = RIGEL_VERSION_STRING;
    obj["compiler"] = RIGEL_CXX_COMPILER_ID;
    obj["compiler_version"] = RIGEL_CXX_COMPILER_VERSION;
    obj["release"] = RIGEL_BUILD_TYPE_STRING;
    obj["build system"] = RIGEL_BUILD_SYSTEM;
    res->set_body(obj.dump());
}

static void set_response_fail(krpc::RestfulResponse*res, int code, const std::string &msg) {
    res->set_status_code(200);
    res->set_content_type("text/plain");
    nlohmann::json obj;
    obj["code"] = code;
    obj["message"] = msg;
    res->set_body(obj.dump());
}

TURBO_FLAG(int32_t, port, 8010, "TCP Port of this server");
TURBO_FLAG(int32_t, idle_timeout_s, -1, "Connection will be closed if there is no "
                                        "read/write operations during the last `idle_timeout_s'");

TURBO_FLAG(std::string, certificate, "", "Certificate file path to enable SSL");
TURBO_FLAG(std::string, private_key, "", "Private key file path to enable SSL");
TURBO_FLAG(std::string, ciphers, "", "Cipher suite used for SSL connections");

int main(int argc, char* argv[]) {

    auto *run_app = turbo::Servlet::instance().run_app();
    run_app->add_option("--port", FLAGS_port, FLAGS_port.help());
    run_app->add_option("--idle_timeout_s", FLAGS_idle_timeout_s, FLAGS_idle_timeout_s.help());
    run_app->add_option("--certificate", FLAGS_certificate, FLAGS_certificate.help());
    run_app->add_option("--private_key", FLAGS_private_key, FLAGS_private_key.help());
    run_app->add_option("--ciphers", FLAGS_ciphers, FLAGS_ciphers.help());
    auto [exit, ret] = turbo::Servlet::instance().run(argc, argv);
    if (exit) {
        return ret;
    }
    // Generally you only need one Server.
    krpc::Server server;

    auto *ins = krpc::RestfulService::instance();

    ins->set_processor("/version", [](const krpc::RestfulRequest* req, krpc::RestfulResponse*res) -> void {
        set_response_version(res);
    });
    ins->set_processor("/create_index", [](const krpc::RestfulRequest* req, krpc::RestfulResponse*res) -> void {
        auto data = nlohmann::json::parse(req->body().to_string());
        rigel::IndexRequest index_request = data.get<rigel::IndexRequest>();
        if (indices.find(index_request.index_name) != indices.end()) {
            set_response_fail(res, static_cast<int>(turbo::StatusCode::kAlreadyExists), "Index already exists");
            return;
        }
        hnswlib::SpaceInterface<float>* space = (index_request.space_type == "IP")
                                                ? static_cast<hnswlib::SpaceInterface<float>*>(new hnswlib::InnerProductSpace(index_request.dimension))
                                                : static_cast<hnswlib::SpaceInterface<float>*>(new hnswlib::L2Space(index_request.dimension));

        auto *index = new hnswlib::HierarchicalNSW<float>(
                space,
                10000,
                index_request.M,
                index_request.ef_construction,
                42,
                true
        );

        indices[index_request.index_name] = index;
        index_settings[index_request.index_name] = data;
        set_response_ok(res);
    });

    ins->set_processor("/load_index", [](const krpc::RestfulRequest* req, krpc::RestfulResponse*res) -> void {
        auto data = nlohmann::json::parse(req->body().to_string());
        std::string index_name = data["index_name"];

        if (indices.find(index_name) != indices.end()) {
            set_response_fail(res, static_cast<int>(turbo::StatusCode::kNotFound), "Index not exists");
            return;
        }

        read_index_from_disk(index_name);
        set_response_ok(res);
    });

    ins->set_processor("/save_index", [](const krpc::RestfulRequest* req, krpc::RestfulResponse*res) -> void {
        auto data = nlohmann::json::parse(req->body().to_string());
        std::string index_name = data["index_name"];

        if (indices.find(index_name) == indices.end()) {
            set_response_fail(res, static_cast<int>(turbo::StatusCode::kNotFound), "Index not exists");
            return;
        }

        write_index_to_disk(index_name);
        set_response_ok(res);
    });

    ins->set_processor("/delete_index", [](const krpc::RestfulRequest* req, krpc::RestfulResponse*res) -> void {
        auto data = nlohmann::json::parse(req->body().to_string());
        std::string index_name = data["index_name"];

        if (indices.find(index_name) == indices.end()) {
            set_response_fail(res, static_cast<int>(turbo::StatusCode::kNotFound), "Index not exists");
            return;
        }

        delete indices[index_name];
        indices.erase(index_name);
        index_settings.erase(index_name);
        set_response_ok(res);
    });

    ins->set_processor("/add_documents", [](const krpc::RestfulRequest* req, krpc::RestfulResponse*res) -> void {
        auto data = nlohmann::json::parse(req->body().to_string());
        rigel::AddDocumentsRequest add_req = data.get<rigel::AddDocumentsRequest>();

        // if missing any keys, return a 400 with the details
        // If vectors is not in the request, return a 400



        if (add_req.ids.size() != add_req.vectors.size()) {
            set_response_fail(res, static_cast<int>(turbo::StatusCode::kInvalidArgument), "Number of IDs does not match number of vectors");
            return;
        }

        if (indices.find(add_req.index_name) == indices.end()) {
            set_response_fail(res, static_cast<int>(turbo::StatusCode::kNotFound), "Index not exists");
            return;
        }

        auto *index = indices[add_req.index_name];

        int current_size = index->cur_element_count;

        if (current_size + add_req.ids.size() > index->max_elements_) {
            index->resizeIndex(current_size + current_size*0.2 + add_req.ids.size());
        }

        for (int i = 0; i < add_req.ids.size(); i++) {
            std::vector<float> vec_data(add_req.vectors[i].begin(), add_req.vectors[i].end());
            index->addPoint(vec_data.data(), add_req.ids[i], 0);
        }
        set_response_ok(res);
    });

    ins->set_processor("/search", [](const krpc::RestfulRequest* req, krpc::RestfulResponse*res) -> void {
        auto data = nlohmann::json::parse(req->body().to_string());
        rigel::SearchRequest search_req = data.get<rigel::SearchRequest>();

        if (indices.find(search_req.index_name) == indices.end()) {
            set_response_fail(res, static_cast<int>(turbo::StatusCode::kNotFound), "Index not exists");
            return;
        }

        auto *index = indices[search_req.index_name];
        std::vector<float> query_vec(search_req.query_vector.begin(), search_req.query_vector.end());
        index->setEf(search_req.ef_search);

        // if (query_vec.size() != thing) {
        //     return crow::response(400, "Query vector size does not match index dimension. Expected " + std::to_string(thing) + " but got " + std::to_string(query_vec.size()));
        // }

        std::priority_queue<std::pair<float, hnswlib::labeltype>> result = index->searchKnn(query_vec.data(), search_req.k);

        nlohmann::json obj;
        std::vector<int> ids;
        std::vector<float> distances;
        while (!result.empty()) {
            ids.push_back(result.top().second);
            distances.push_back(result.top().first);
            result.pop();
        }

        std::reverse(ids.begin(), ids.end());
        std::reverse(distances.begin(), distances.end());

        obj["hits"] = ids;
        obj["distances"] = distances;

        res->set_status_code(200);
        res->set_content_type("text/plain");
        res->set_body(obj.dump());
    });


    // Start the server.
    krpc::ServerOptions options;
    options.idle_timeout_sec = turbo::get_flag(FLAGS_idle_timeout_s);
    if(!turbo::get_flag(FLAGS_certificate).empty() && !turbo::get_flag(FLAGS_private_key).empty()) {
        options.mutable_ssl_options()->default_cert.certificate = turbo::get_flag(FLAGS_certificate);
        options.mutable_ssl_options()->default_cert.private_key = turbo::get_flag(FLAGS_private_key);
        options.mutable_ssl_options()->ciphers = turbo::get_flag(FLAGS_ciphers);
    }
    if (server.Start(turbo::get_flag(FLAGS_port), &options) != 0) {
        LOG(ERROR) << "Fail to start HttpServer";
        return -1;
    }

    // Wait until Ctrl-C is pressed, then Stop() and Join() the server.
    server.RunUntilAskedToQuit();
    return 0;

}